syntax = "proto3";

package tensorflow.data;

import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/protobuf/data_service.proto";
import "tensorflow/core/protobuf/snapshot.proto";
import "tensorflow/tsl/protobuf/status.proto";

// Next tag: 2
message DatasetDef {
  // We represent datasets as tensorflow GraphDefs which define the operations
  // needed to create a tf.data dataset.
  GraphDef graph = 1;
}

// Next tag: 3
message IterationKeyDef {
  string name = 1;
  int64 iteration = 2;
}

// Next tag: 14
message TaskDef {
  reserved 6;
  // The dataset to iterate over.
  oneof dataset {
    DatasetDef dataset_def = 1;
    string path = 2;
  }
  string dataset_id = 3;
  int64 task_id = 4;
  int64 iteration_id = 5;
  // In distributed epoch processing mode, we use one split provider for each
  // source that feeds into the dataset. In parallel_epochs mode,
  // `num_split_providers` is always zero.
  int64 num_split_providers = 9;
  // Address of the worker that the task is assigned to.
  string worker_address = 8;
  ProcessingModeDef processing_mode_def = 10;
  // Optional number of consumers. If set, the results of the task will be
  // provided to consumers round-robin.
  oneof optional_num_consumers {
    int64 num_consumers = 7;
  }
  // Number of workers and the worker index. These are only populated when the
  // `processing_mode_def` specifies a static sharding policy.
  int64 num_workers = 11;
  int64 worker_index = 12;
  // True if cross-trainer cache is enabled.
  bool use_cross_trainer_cache = 13;
}

// Next tag: 9
message TaskInfo {
  // The address of the worker processing the task.
  string worker_address = 1;
  // The data transfer servers of the worker processing the task.
  repeated DataTransferServerInfo transfer_servers = 8;
  // Tags attached to the worker. This allows reading from selected workers.
  // For example, by applying a "COLOCATED" tag, tf.data service is able to read
  // from the local tf.data worker if one exists, then from off-TF-host workers,
  // to avoid cross-TF-host reads.
  repeated string worker_tags = 6;
  // The task id.
  int64 task_id = 2;
  // The id of the iteration that the task is part of.
  int64 iteration_id = 3;
  // The UID of the worker Borg job, used for telemetry.
  int64 worker_uid = 7;
  // The round to start reading from the task in. For non-round-robin reads,
  // this is always 0.
  int64 starting_round = 5;
  reserved 4;
}

// Next tag: 5
message SnapshotTaskDef {
  // The base directory at which the snapshot is being materialized.
  string base_path = 1;
  // The index of the stream that the worker has been assigned to process.
  int64 stream_index = 2;
  // The number of source datasets (split providers).
  int64 num_sources = 3;
  // Snapshot metadata including the element spec and compression method.
  experimental.DistributedSnapshotMetadata metadata = 4;
}

// Next tag: 5
message SnapshotTaskProgress {
  SnapshotTaskDef snapshot_task = 1;
  // True if the snapshot is complete successfully. Unset if the snapshot is not
  // complete or an error has occurred.
  bool completed = 2;
  // If any error occurs during the snapshot processing, `status` will be filled
  // with the error status.
  StatusProto status = 3;
}

message SnapshotStreamInfo {
  // The index of the stream being processed or having been processed.
  int64 index = 1;

  enum State {
    // Unspecified.  Invalid state.
    UNSPECIFIED = 0;
    // The dispatcher thinks the stream has a live worker.
    ASSIGNED = 1;
    // The dispatcher doesn't think the stream has a live worker.
    ORPHAN = 2;
    // The dispatcher doesn't know whether or not the stream has a live worker.
    UNKNOWN = 3;
    // The dispatcher thinks the stream has no more splits left to be processed.
    DONE = 4;
  }
  State state = 2;
}

// Specifies which tf.data service workers to read from.
enum TargetWorkers {
  TARGET_WORKERS_UNSPECIFIED = 0;
  // tf.data service runtime decides which workers to read from.
  TARGET_WORKERS_AUTO = 1;
  // Reads from any available worker.
  TARGET_WORKERS_ANY = 2;
  // Only reads from local workers. If no local worker is found, it is an error.
  TARGET_WORKERS_LOCAL = 3;
}

// Information about one of a worker server's data transfer servers.
message DataTransferServerInfo {
  string protocol = 1;
  string address = 2;
  // If provided, properties of the server used to determine compatibility with
  // a client.
  bytes compatibility_info = 3;
}
